#!/usr/bin/env python

import simsym
from symtypes import *
import importlib
import z3
import testgen
import z3printer
import argparse
import z3util

parser = argparse.ArgumentParser()
parser.add_argument('module', metavar='MODULE', default='fs', action='store',
                    help='Module to test (e.g., models.fs)')
parser.add_argument('-t', '--test-file',
                    help='Test generator output file')

def simple_parse_functions(module):
    base = module.model_class
    callnames = {name for name in dir(base)
                 if getattr(getattr(base, name), 'is_model_function', False)}

    return callnames

def simple_test(base, *calls):

    # Create arguments for each call.  We reuse these arguments in
    # each permutation, so each call receives the same arguments each
    # time we test it.
    args = []
    for call in calls:
        arg_name = '%s' % (call.__name__)
        args.append(call.arg_struct_type.var(arg_name))

    init = base.var(base.__name__)

    # Build the Python arguments dictionary and copy each
    # argument, just in case the call mutates it
    arg_struct = args[0]
    cargs = {arg: getattr(arg_struct, arg).copy()
                     for arg in arg_struct._fields}
    # Invoke the call
    nstate = init.copy()

    try:
        res = call(nstate, **cargs)
    except BaseException, e:
        print str(e)
    finally:
        #model.cur_thread_idx = None
        note(('end', 0, res))
    return res

def simple_do_call(base,call,test_writer):
    test_writer.begin_call_set([call])
    condlists = []

    for sar in symbolic_apply(simple_test, base, call):
        if sar.type == 'value':
            condlists.append(sar.path_condition)
        test_writer.on_result(sar)

    test_writer.end_call_set()

    print "----------------------------------"
    print call.__name__, ":: exec path is ", len(condlists)

    i=0
    for cond in condlists:
        i+=1
        print " ",i,"--",cond

    return call.__name__, len(condlists)
#######################################

class IsomorphicMatch(object):
    """Construct an expression that matches isomorphisms of a set of conditions.

    By using the negation of the constructed expression, we can
    enumerate non-isomorphic models."""

    ## XXX handling FDs and timestamps might be better done by treating
    ## them as supporting order, rather than supporting just equality;
    ## the isomorphism condition would be the values being in the same
    ## order, rather than in the same equality pattern.

    def __init__(self):
        # For conditions subject to equality isomorphism, this maps
        # from realms (specifically, Interpreter instances) to
        # equivalence classes of expressions.  Each equivalence class
        # contains the set of known expressions that evaluate to the
        # same literal value.  Specifically,
        #   {Interpreter: RepMap}
        #   RepMap := {HashableAst literal (e.g., 2 or f!0): Reps}
        #   Reps   := AstSet of expressions that evaluate to the representative
        self.__repmaps = collections.defaultdict(
            lambda: collections.defaultdict(z3util.AstSet))
        # Set of isomorphism conditions for value isomorphisms.
        self.__conds = z3util.AstSet()

    def add(self, realm, expr, val, env):
        """Add a condition on the value of expr.

        expr and val must be instances of simsym.Symbolic.  env must
        be a simsym.SymbolicApplyResult that provides the context for
        interpreting constants in expr and val.

        For most expr's, this will require that an isomorphic
        assignment assign val to expr.

        If realm is an Interpreter, then this will require that all
        keys in the Interpreter be distinct, but it will not place
        constraints on their specific values.  This makes it possible
        to match isomorphic "equality patterns" between values of
        uninterpreted sorts, or values we wish to treat as
        uninterpreted.
        """

        if not isinstance(expr, simsym.Symbolic):
            raise TypeError("Expected instance of simsym.Symbolic, got %r" % expr)
        if not isinstance(val, simsym.Symbolic):
            raise TypeError("Expected instance of simsym.Symbolic, got %r" % val)

        if realm is None:
            # Default realm.  Use value equality.
            if not isinstance(expr, simsym.SBool):
                print 'WARNING: Interpreted sort assignment:', \
                    type(expr), expr, val
            self.__conds.add(expr == val)
        elif isinstance(realm, testgen.Interpreter):
            # Use equality isomorphism within this realm
            hval = z3util.HashableAst(val)
            self.__repmaps[realm][z3util.HashableAst(val)].add(expr)
        else:
            raise ValueError("Unknown realm type %r" % realm)

    def condition(self):
        """Return the isomorphism condition."""
        conds = list(self.__conds)

        for rep_map in self.__repmaps.itervalues():
            distinct = []
            for reps in rep_map.itervalues():
                reps = list(reps)
                # Require each representative group to be distinct
                distinct.append(reps[0])
                # Require all expressions within the representative
                # group to be equal
                conds.append(simsym.symeq(*reps))
            if len(distinct) > 1:
                conds.append(simsym.distinct(*distinct))

        return simsym.symand(conds)


def expr_vars(e):
    """Return an AstSet of uninterpreted constants in e.

    Uninterpreted constants are what people normally think of as
    "variables".  This is in contrast with interpreted constants such
    as the number 2.  Note that this will also return values that
    belong to universes of uninterpreted sorts, since there is no
    distinguishable difference between these and other uninterpreted
    constants.
    """

    res = z3util.AstSet()
    def rec(e):
        if not z3.is_ast(e):
            return
        if z3.is_const(e) and e.decl().kind() == z3.Z3_OP_UNINTERPRETED:
            res.add(e)
            return
        for child in e.children():
            rec(child)
    rec(simsym.unwrap(e))
    return res


class TestWriter(object):
    def __init__(self, test_file, testgen):
        self.test_file = test_file
        if test_file and testgen:
            self.testgen = testgen(test_file)
        else:
            self.testgen = None

    def begin_call_set(self, callset):

        self.callset = callset
        self.npath = self.nmodel = self.nerror = 0

        if self.testgen:
            self.testgen.begin_call_set(callset)

    def keep_going(self):
        return True
        #TODO: add args
        #return self.nmodel < args.max_testcases

    def on_result(self, result):
        self.npath += 1

        if result.type == 'exception':
            self.nerror += 1
            return

        if not self.testgen:
            return

        e = result.path_condition

        ## This can potentially reduce the number of test cases
        ## by, e.g., eliminating irrelevant variables from e.
        ## The effect doesn't seem significant: one version of Fs
        ## produces 3204 test cases without simplify, and 3182 with.
        e = simsym.simplify(e)

        # Find the uninterpreted constants in the path condition.  We
        # omit assumptions because uninterpreted constants that appear
        # only in assumptions generally don't represent that the model
        # actually "touched".  We use the simplified expression
        # because the final state comparison in original expression
        # contains a lot of trivial expressions like x==x for all
        # state variables x, and we don't care about these
        # uninterpreted constants.
        e_vars = expr_vars(
            simsym.simplify(
                simsym.symand(
                    result.get_path_condition_list(
                        with_assume=False, with_det=True))))

        if self.testgen:
            self.testgen.begin_path(result)

        self.npathmodel = 0
        while self.keep_going():
            # XXX Would it be faster to reuse the solver?
            check = simsym.check(e)
            if check.is_unsat: break
            if check.is_unknown:
                # raise Exception('Cannot enumerate: %s' % str(e))
                print 'Cannot enumerate, moving on..'
                print 'Failure reason:', check.reason
                break

            if 'array-ext' in check.z3_model.sexpr():
                # Work around some non-deterministic bug that causes
                # Z3 to occasionally produce models containing
                # 'array-ext' applications that break evaluation.
                print 'Warning: Working around array-ext bug'
                for i in range(10):
                    check = simsym.check(e)
                    if 'array-ext' not in check.z3_model.sexpr():
                        break
                else:
                    print 'Workaround failed; this won\' end well'

            testid = ('_'.join(c.__name__ for c in self.callset) +
                      '_' + result.pathid + '_' + str(self.npathmodel))

            assignments = self.__on_model(result, check.z3_model, e, testid)
            if assignments is None:
                break

            # Construct the isomorphism condition for the assignments
            # used by testgen.  This tells us exactly what values
            # actually mattered to test case generation.  However,
            # this set isn't perfect: testgen may have queried
            # assignments that didn't actually matter to the
            # function's behavior (e.g., flags that didn't matter
            # because the function will return an error anyway, etc).
            # To filter out such uninterpreted constants, we only
            # consider those that were *both* used in an assignment by
            # testgen and appeared in the path condition expression.
            # XXX We should revisit this and see how much difference
            # this makes.
            same = IsomorphicMatch()
            for realm, rassigns in assignments.iteritems():
                for aexpr, val in rassigns:
                    aexpr_vars = expr_vars(aexpr)
                    if not aexpr_vars.isdisjoint(e_vars):
                        same.add(realm, aexpr, val, result)

            isocond = same.condition()

            # Construct constraint for next test
            notsame = simsym.symnot(isocond)
            e = simsym.symand([e, notsame])

        if self.testgen:
            self.testgen.end_path()

    def __on_model(self, result, model, constraint, testid):
        self.nmodel += 1
        res = None

        if self.testgen:
            smodel = result.get_model(model)
            smodel.track_assignments(True)
            self.testgen.on_model(testid, smodel, constraint)
            res = smodel.assignments()

        self.npathmodel += 1
        return res

    def end_call_set(self):
        if self.testgen:
            self.testgen.end_call_set()

    def finish(self):
        if self.testgen:
            self.testgen.finish()


max_path=0
cp_list=[]

def main(spec_args):
    global args
    args = spec_args
    z3printer._PP.max_lines = float('inf')
    m = importlib.import_module(args.module)

    testgen = m.model_testgen if hasattr(m, 'model_testgen') else None
    if testgen is None and args.test_file:
        parser.error("No test case generator for this module")

    test_writer = TestWriter(args.test_file, testgen)

    max_path_num = 0
    callnames = simple_parse_functions(m)
    callset = [getattr(m.model_class, callname) for callname in callnames]
    for call in callset:
        call_name, path_num=simple_do_call(m.model_class, call, test_writer)
        cp_list.append((call_name,path_num))
        if max_path_num<path_num:
            call_name_max,max_path_num=call_name,path_num

    test_writer.finish()

    print "======call name, path num========="
    print cp_list
    print call_name_max," :: max path num is ",max_path_num

if __name__ == "__main__":
    main(parser.parse_args())
